import sys
import os
import cv2
import numpy as np
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout,
    QHBoxLayout, QPushButton, QLabel, QLineEdit,
    QFileDialog, QTextEdit, QComboBox
)
from PySide6.QtGui import QPixmap, QImage
from PySide6.QtCore import Qt


# ------------------ 解决 OpenCV 中文路径问题 ------------------
def imread_unicode(path, flags=cv2.IMREAD_COLOR):
    """支持中文路径的 imread"""
    if not os.path.exists(path):
        return None
    data = np.fromfile(path, dtype=np.uint8)
    return cv2.imdecode(data, flags)


def imwrite_unicode(path, img):
    """支持中文路径的 imwrite"""
    ext = os.path.splitext(path)[1]
    success, buf = cv2.imencode(ext, img)
    if success:
        buf.tofile(path)
    return success
# -------------------------------------------------------------


class PCBImageProcessor:
    def __init__(self, scheme="Blue"):
        # 默认颜色表
        self.baseColors = {
            "DGreen": "#193522",   # 深绿
            "Green": "#F9E195",   # 浅绿
            "Black": "#061008",   # 黑
            "White": "#E6EAEB",   # 白
        }

        # 可选的替换组
        self.colorSchemes = {
            "Blue": {
                "DBlue": "#161F7D",
                "Blue": "#5DA7E3",
            },
            "Red": {
                "DRed": "#7D1616",
                "Red": "#E35D5D",
            },
            "Purple": {
                "DPurple": "#3E166C",
                "Purple": "#B85DE3",
            },
            "Yellow": {
                "DYellow": "#7D6F16",
                "Yellow": "#E3D65D",
            }
        }

        self.set_scheme(scheme)

    def set_scheme(self, scheme):
        """设置当前配色方案"""
        if scheme not in self.colorSchemes:
            scheme = "Blue"
        self.hexColors = {**self.baseColors, **self.colorSchemes[scheme]}
        self.colors = {name: self.hex_to_BGR(value) for name, value in self.hexColors.items()}

    def hex_to_BGR(self, hex_color):
        hex_color = hex_color.lstrip('#')
        return np.array([int(hex_color[i:i + 2], 16) for i in (4, 2, 0)])

    def simplify_colors(self, input_path, output_dir="./1.颜色简化结果/"):
        image = imread_unicode(input_path)
        if image is None:
            raise FileNotFoundError(f"无法读取图片: {input_path}")

        scale_factor = 0.25
        image = cv2.resize(image, (int(image.shape[1] * scale_factor),
                                   int(image.shape[0] * scale_factor)),
                           interpolation=cv2.INTER_LINEAR)

        # 使用 numpy 广播简化颜色
        reshaped = image.reshape((-1, 3))
        color_values = np.array(list(self.colors.values()))
        distances = np.linalg.norm(reshaped[:, None] - color_values[None, :], axis=2)
        nearest_color_idx = np.argmin(distances, axis=1)
        simplified = color_values[nearest_color_idx].reshape(image.shape)

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        out_path = os.path.join(output_dir, os.path.splitext(os.path.basename(input_path))[0] + "_simplified.png")
        imwrite_unicode(out_path, simplified)
        return out_path

    def separate_colors(self, simplified_path, output_dir="./2.颜色分离结果/"):
        image = imread_unicode(simplified_path)
        if image is None:
            raise FileNotFoundError(f"无法读取图片: {simplified_path}")

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        for name, hex_color in self.hexColors.items():
            bgr_color = self.hex_to_BGR(hex_color)
            mask = cv2.inRange(image, bgr_color, bgr_color)
            mask_image = cv2.bitwise_and(image, image, mask=mask)
            imwrite_unicode(os.path.join(output_dir, f"{name}.png"), mask_image)

        return output_dir

    def combine_colors(self, separated_dir="./2.颜色分离结果/", output_dir="./3.颜色合并结果/"):
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        def load_img(name):
            return imread_unicode(os.path.join(separated_dir, f"{name}.png"))

        # 判断当前使用的方案，自动获取组名
        scheme_keys = [k for k in self.hexColors.keys() if k not in self.baseColors]
        if len(scheme_keys) >= 2:
            deep, light = scheme_keys[:2]
        else:
            deep, light = None, None

        # 顶层阻焊层：深色 + 浅色
        if deep and light:
            top_mask = cv2.bitwise_or(load_img(deep), load_img(light)) # type: ignore
            imwrite_unicode(os.path.join(output_dir, "顶层阻焊层.png"), top_mask)

            # 顶层铜皮层：浅色 + 黑
            top_copper = cv2.bitwise_or(load_img(light), load_img("Black")) # type: ignore
            imwrite_unicode(os.path.join(output_dir, "顶层铜皮层.png"), top_copper)

        # 顶层丝印层：白色
        imwrite_unicode(os.path.join(output_dir, "顶层丝印层.png"), load_img("White"))

        # 底层铜皮层：深绿
        imwrite_unicode(os.path.join(output_dir, "底层铜皮层.png"), load_img("DGreen"))

        # 底层阻焊层：浅绿
        imwrite_unicode(os.path.join(output_dir, "底层阻焊层.png"), load_img("Green"))

        return output_dir


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.processor = PCBImageProcessor()
        self.setWindowTitle("PCB 图片处理软件")

        # 中央控件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)

        # 输入框 + 浏览按钮
        input_layout = QHBoxLayout()
        self.input_edit = QLineEdit()
        self.input_edit.setPlaceholderText("输入图片路径或点击浏览...")
        browse_btn = QPushButton("浏览")
        browse_btn.clicked.connect(self.browse_file)
        input_layout.addWidget(self.input_edit)
        input_layout.addWidget(browse_btn)
        layout.addLayout(input_layout)

        # 颜色方案下拉框
        scheme_layout = QHBoxLayout()
        scheme_label = QLabel("PCB颜色方案:")
        self.scheme_combo = QComboBox()
        self.scheme_combo.addItems(["Blue", "Red", "Purple", "Yellow"])
        self.scheme_combo.currentTextChanged.connect(self.change_scheme)
        scheme_layout.addWidget(scheme_label)
        scheme_layout.addWidget(self.scheme_combo)
        layout.addLayout(scheme_layout)

        # 功能按钮
        btn_layout = QHBoxLayout()
        simplify_btn = QPushButton("颜色简化")
        separate_btn = QPushButton("颜色分离")
        combine_btn = QPushButton("颜色合成")
        simplify_btn.clicked.connect(self.run_simplify)
        separate_btn.clicked.connect(self.run_separate)
        combine_btn.clicked.connect(self.run_combine)
        btn_layout.addWidget(simplify_btn)
        btn_layout.addWidget(separate_btn)
        btn_layout.addWidget(combine_btn)
        layout.addLayout(btn_layout)

        # 预览区
        self.image_label = QLabel("预览区")
        self.image_label.setAlignment(Qt.AlignCenter)  # type: ignore
        layout.addWidget(self.image_label)

        # 状态区
        self.log = QTextEdit()
        self.log.setReadOnly(True)
        layout.addWidget(self.log)

    def browse_file(self):
        file, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.bmp)")
        if file:
            self.input_edit.setText(file)
            self.show_image(file)

    def show_image(self, path):
        image = QImage(path)
        if image.isNull():
            return
        pixmap = QPixmap.fromImage(image).scaled(400, 300, Qt.KeepAspectRatio)  # type: ignore
        self.image_label.setPixmap(pixmap)

    def log_message(self, msg):
        self.log.append(msg)

    def change_scheme(self, scheme):
        self.processor.set_scheme(scheme)
        self.log_message(f"切换颜色方案: {scheme}")

    def run_simplify(self):
        path = self.input_edit.text().strip()
        if not path:
            self.log_message("请先输入图片路径")
            return
        try:
            out = self.processor.simplify_colors(path)
            self.log_message(f"颜色简化完成: {out}")
            self.show_image(out)
        except Exception as e:
            self.log_message(str(e))

    def run_separate(self):
        path = self.input_edit.text().strip()
        base_name = os.path.splitext(os.path.basename(path))[0]
        simplified_path = f"./1.颜色简化结果/{base_name}_simplified.png"
        try:
            out_dir = self.processor.separate_colors(simplified_path)
            self.log_message(f"颜色分离完成: {out_dir}")
        except Exception as e:
            self.log_message(str(e))

    def run_combine(self):
        try:
            out_dir = self.processor.combine_colors()
            self.log_message(f"颜色合成完成: {out_dir}")
        except Exception as e:
            self.log_message(str(e))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    win = MainWindow()
    win.resize(800, 600)
    win.show()
    sys.exit(app.exec())
